import cv2
import numpy as np
import tensorflow as tf
import keras.backend as K
import matplotlib as mpl
import matplotlib.pyplot as plt
from tqdm import tqdm
from keras.models import Sequential
from keras.callbacks import LearningRateScheduler
from keras.preprocessing.image import img_to_array
from keras.layers import Dense,Conv2D,MaxPool2D,Flatten,Dropout



# hyper variables sets
epoch = 100
pic_size = 32, 32
model_version = "1M100E32X_1"
max_catory = 999999999

# load train files
train_x = []
train_y = []
train_file = open("./train.csv").read()
train_data = train_file.split("\n")[1:]
cator1 = cator2 = cator3 = cator4 = file_not_found = 0
for i in tqdm(train_data):
    args = i.split(",")
    args[1] = int(args[1])

    if args[1] <= 6 and cator1 <= max_catory: 
        cator = [1,0,0,0]
        cator1 += 1
    elif args[1] > 6 and args[1] <= 14 and cator2 <= max_catory: 
        cator = [0,1,0,0]
        cator2 += 1
    elif args[1] > 14 and args[1] <= 37 and cator3 <= max_catory: 
        cator = [0,0,1,0]
        cator3 += 1
    elif cator4 <= max_catory : 
        cator = [0,0,0,1]
        cator4 += 1
    else: pass
    if cator1 >= max_catory \
        and cator2 >= max_catory \
        and cator3 >= max_catory \
        and cator4 >= max_catory: break

    try:
        train_x.append(img_to_array(cv2.resize(cv2.imread(args[0],0),pic_size)))
        train_y.append(cator)
    except Exception:
        file_not_found += 1
        continue


print(cator1,cator2,cator3,cator4,"Missing Files: "+str(file_not_found))

train_x = np.array(train_x)
train_y = np.array(train_y)

# load test files
test_x = []
test_y = []
test_file = open("./val.csv").read()
test_data = test_file.split("\n")[1:]
#test_data = random.sample(test_data,8000)
cator1 = cator2 = cator3 = cator4 = file_not_found = 0
for i in tqdm(test_data):
    args = i.split(",")
    args[1] = int(args[1])
    if args[1] <= 6 and cator1 <= max_catory: 
        cator = [1,0,0,0]
        cator1 += 1
    elif args[1] > 6 and args[1] <= 14 and cator2 <= max_catory: 
        cator = [0,1,0,0]
        cator2 += 1
    elif args[1] > 14 and args[1] <= 37 and cator3 <= max_catory: 
        cator = [0,0,1,0]
        cator3 += 1
    elif cator4 <= max_catory : 
        cator = [0,0,0,1]
        cator4 += 1
    else: pass
    if cator1 >= max_catory \
        and cator2 >= max_catory \
        and cator3 >= max_catory \
        and cator4 >= max_catory: break

    try:
        test_x.append(img_to_array(cv2.resize(cv2.imread(args[0],0),pic_size)))
        test_y.append(cator)
    except Exception:
        file_not_found += 1
        continue

print(cator1,cator2,cator3,cator4,"Missing Files: "+str(file_not_found))

test_x = np.array(test_x)
test_y = np.array(test_y)

# create model
model = Sequential()

model.add(
    Conv2D(64, kernel_size=[5,5],strides=2,
           activation="relu",input_shape=train_x.shape[1:])
)
model.add(
    MaxPool2D(pool_size=[3,3],strides=2)
)

model.add(Flatten())
#model.add(Dropout(0.1))
for i in range(2):
    model.add(Dense(256,activation="relu"))
model.add(
    Dense(4,activation="softmax")
)


# run model
model.compile(
    loss="categorical_crossentropy",optimizer="adam",metrics=["accuracy"]
)

def scheduler(lepoch):
    if lepoch % 20 == 0 and lepoch > 2:
        lr = K.get_value(model.optimizer.lr)
        K.set_value(model.optimizer.lr, lr * 0.05)
    return K.get_value(model.optimizer.lr)


lr_new = LearningRateScheduler(scheduler)


history = model.fit(
    train_x,
    train_y,
    epochs=epoch,
    validation_data=(test_x,test_y),
    workers=6,
    use_multiprocessing=True,
    callbacks=[lr_new]
)

model.save("LitterCollecter_{}.keras".format(model_version))

# plot history
mpl.rcParams['font.family'] = 'SimHei'
mpl.rcParams['axes.unicode_minus'] = False

plt.plot(history.history['accuracy'])
plt.plot(history.history['val_accuracy'])
plt.title('模型准确率')
plt.ylabel('准确率')
plt.xlabel('训练次数')
plt.legend(['训练集', '验证集'], loc='upper left')
plt.show()

plt.plot(history.history['loss'])
plt.plot(history.history['val_loss'])
plt.title('模型损失')
plt.ylabel('损失')
plt.xlabel('训练次数')
plt.legend(['训练集', '验证集'], loc='upper left')
plt.show()


pre_x = img_to_array(
    cv2.resize(
        cv2.imread("F:/WorkFlow Corporation/Projects/AITest/files/val/label23/img_11525.jpg",0),
        pic_size
    )
)
pre_x = tf.expand_dims(pre_x,0)
print(list(model.predict(pre_x)))
input()
